import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error, mean_absolute_error

"""
特征值（输入量）：历史奖牌数据、是否为主办国，每个国家每年参加运动员的总数
标签值（输出量）：金牌数

随机森林1.0 输出两个图：
1. 实际值与预测值的对比图
   - 作用：展示模型预测的金牌数量与真实金牌数量的匹配程度。如果点分布接近红色理想预测线（y=x），说明预测性能较好。
2. 预测误差的分布直方图
   - 作用：揭示模型预测误差的分布情况。如果误差集中在0附近，说明模型整体误差较小；误差的分布趋势和范围可以帮助判断模型的稳健性。
"""

# 国家关系映射表，medal_counts和athletes两个文件表示不同
noc_mapping = {
    "United States": "USA",
    "Greece": "GRE",
    "Germany": "DEU",
    "France": "FRA",
    "Great Britain": "GBR",  # 映射为 United Kingdom
    "China": "CHN",
    "Denmark": "DNK",
    "Netherlands": "NED",
    "Hungary": "HUN",
    "Austria": "AUT",
    "Australia": "AUS",
    "Sweden": "SWE",
    "Italy": "ITA",
    "Belgium": "BEL",
    "Switzerland": "CHE",
    "Canada": "CAN",
    "Norway": "NOR",
    "Russia": "RUS",
    "Japan": "JPN",
    "Finland": "FIN",
    "Poland": "POL",
    "Turkey": "TUR",
    "Brazil": "BRA",
    "Romania": "ROU",
    "Czechoslovakia": "CZE",  # 已映射为 CZE
    "Czech Republic": "CZE",
    "South Korea": "KOR",
    "Spain": "ESP",
    "Argentina": "ARG",
    "New Zealand": "NZL",
    "South Africa": "RSA",
    "Mexico": "MEX",
    "Iran": "IRN",
    "Ethiopia": "ETH",
    "Kenya": "KEN",
    "Ukraine": "UKR",
    "Cuba": "CUB",
    "Portugal": "PRT",
    "Azerbaijan": "AZE",
    "Bahrain": "BAH",
    "Belarus": "BLR",
    "Bulgaria": "BUL",
    "Chile": "CHL",
    "Colombia": "COL",
    "Dominican Republic": "DOM",
    "Ecuador": "ECU",
    "Egypt": "EGY",
    "Estonia": "EST",
    "Georgia": "GEO",
    "Grenada": "GRD",
    "Guatemala": "GTM",
    "Hong Kong": "HKG",
    "India": "IND",
    "Indonesia": "IDN",
    "Israel": "ISR",
    "Jamaica": "JAM",
    "Jordan": "JOR",
    "Kazakhstan": "KAZ",
    "Kosovo": "KOS",
    "Kyrgyzstan": "KGZ",
    "Latvia": "LVA",
    "Lithuania": "LTU",
    "Moldova": "MDA",
    "Mongolia": "MNG",
    "Morocco": "MAR",
    "Namibia": "NAM",
    "North Korea": "PRK",
    "North Macedonia": "MKD",  # 映射为 North Macedonia
    "Norway": "NOR",
    "Pakistan": "PAK",
    "Panama": "PAN",
    "Paraguay": "PRY",
    "Peru": "PER",
    "Philippines": "PHL",
    "Qatar": "QAT",
    "Refugee Olympic Team": "ROT",
    "San Marino": "SMR",
    "Saudi Arabia": "KSA",
    "Serbia": "SRB",  # 映射为 Serbia
    "Singapore": "SGP",
    "Slovakia": "SVK",
    "Slovenia": "SVN",
    "Syria": "SYR",
    "Tajikistan": "TJK",
    "Thailand": "THA",
    "Trinidad and Tobago": "TTO",
    "Turkmenistan": "TKM",
    "Tuvalu": "TUV",
    "Uganda": "UGA",
    "United Arab Emirates": "UAE",
    "Uruguay": "URY",
    "Uzbekistan": "UZB",
    "Vietnam": "VNM",
    "Yemen": "YEM",
    "Zambia": "ZMB",
    "Zimbabwe": "ZWE",
    "Antigua and Barbuda": "ATG",
    "Armenia": "ARM",
    "Bahamas": "BHS",
    "Bangladesh": "BGD",
    "Barbados": "BRB",
    "Benin": "BEN",
    "Bermuda": "BMU",
    "Bhutan": "BTN",
    "Bolivia": "BOL",
    "Bosnia and Herzegovina": "BIH",
    "Botswana": "BWA",
    "Brunei": "BRN",
    "Burkina Faso": "BFA",
    "Burundi": "BDI",
    "Cambodia": "KHM",
    "Cameroon": "CMR",
    "Cape Verde": "CPV",
    "Central African Republic": "CAF",
    "Chad": "TCD",
    "Comoros": "COM",
    "Congo (Brazzaville)": "COG",
    "Congo (Kinshasa)": "COD",
    "Costa Rica": "CRI",
    "Côte d'Ivoire": "CIV",
    "Croatia": "HRV",
    "Cyprus": "CYP",
    "Djibouti": "DJI",
    "Dominica": "DMA",
    "El Salvador": "SLV",
    "Equatorial Guinea": "GNQ",
    "Eritrea": "ERI",
    "Eswatini": "SWZ",  # 映射为 Swaziland
    "Fiji": "FJI",
    "Gabon": "GAB",
    "Gambia": "GMB",
    "Ghana": "GHA",
    "Grenada": "GRD",
    "Guinea": "GIN",
    "Guinea-Bissau": "GNB",
    "Guyana": "GUY",
    "Haiti": "HTI",
    "Honduras": "HND",
    "Iceland": "ISL",
    "Ireland": "IRL",
    "Kazakhstan": "KAZ",
    "Kenya": "KEN",
    "Kiribati": "KIR",
    "Kuwait": "KWT",
    "Kyrgyzstan": "KGZ",
    "Laos": "LAO",
    "Lebanon": "LBN",
    "Lesotho": "LSO",
    "Liberia": "LBR",
    "Libya": "LBY",
    "Liechtenstein": "LIE",
    "Luxembourg": "LUX",
    "Madagascar": "MDG",
    "Malawi": "MWI",
    "Malaysia": "MYS",
    "Maldives": "MDV",
    "Mali": "MLI",
    "Malta": "MLT",
    "Marshall Islands": "MHL",
    "Mauritania": "MRT",
    "Mauritius": "MUS",
    "Micronesia": "FSM",
    "Monaco": "MCO",
    "Montenegro": "MNE",
    "Mozambique": "MOZ",
    "Myanmar": "MMR",
    "Nauru": "NRU",
    "Nepal": "NPL",
    "Nicaragua": "NIC",
    "Niger": "NER",
    "Nigeria": "NGA",
    "Oman": "OMN",
    "Palau": "PLW",
    "Papua New Guinea": "PNG",
    "Qatar": "QAT",
    "Romania": "ROU",
    "Russia": "RUS",
    "Rwanda": "RWA",
    "Saint Kitts and Nevis": "KNA",
    "Saint Lucia": "LCA",
    "Saint Vincent and the Grenadines": "VCT",
    "Samoa": "WSM",
    "San Marino": "SMR",
    "Sao Tome and Principe": "STP",
    "Saudi Arabia": "KSA",
    "Senegal": "SEN",
    "Serbia": "SRB",  # 映射为 Serbia
    "Seychelles": "SYC",
    "Sierra Leone": "SLE",
    "Singapore": "SGP",
    "Slovakia": "SVK",
    "Slovenia": "SVN",
    "Solomon Islands": "SLB",
    "Somalia": "SOM",
    "South Africa": "ZAF",
    "South Korea": "KOR",
    "South Sudan": "SSD",
    "Sri Lanka": "LKA",
    "Sudan": "SDN",
    "Suriname": "SUR",
    "Sweden": "SWE",
    "Swaziland": "SWZ",  # 映射为 Eswatini
    "Switzerland": "CHE",
    "Syria": "SYR",
    "Taiwan": "TPE",  # 与noc_mapping一致
    "Tajikistan": "TJK",
    "Tanzania": "TZA",
    "Thailand": "THA",
    "Timor-Leste": "TLS",
    "Togo": "TGO",
    "Tonga": "TON",
    "Tunisia": "TUN",
    "Turkey": "TUR",
    "Turkmenistan": "TKM",
    "Tuvalu": "TUV",
    "Uganda": "UGA",
    "United Kingdom": "GBR",
    "United States": "USA",
    "Uruguay": "URY",
    "Uzbekistan": "UZB",
    "Vanuatu": "VUT",
    "Venezuela": "VEN",
    "Vietnam": "VNM",
    "Yemen": "YEM",
    "Zambia": "ZMB",
    "Zimbabwe": "ZWE"
}

# 加载数据集  这些文件是经过预处理的，包括国家关系名称映射处理，programs中缺失值处理
medal_counts = pd.read_csv('2025_Problem_C/processed1_medal_counts.csv', encoding='latin1')
hosts = pd.read_csv('Processed/processed_hosts.csv', encoding='latin1')
programs = pd.read_csv('2025_Problem_C/processed1_programs.csv', encoding='latin1')
athletes = pd.read_csv('Processed/processed_athletes.csv', encoding='latin1')

# 假设你的数据已经加载到 'athletes' 数据框中
# 删除 'Year' 列中小于 1906 的所有行
athletes = athletes[athletes['Year'] >= 1906]

# 数据预处理 对host文件进行处理，解决了乱码问题
hosts.columns = hosts.columns.str.strip()  # 清理列名中的多余空格
hosts.rename(columns={"ï»¿Year": "Year"}, inplace=True)  # 修复可能的编码问题
hosts['Host'] = hosts['Host'].str.replace('Â', '').str.strip()  # 清理特殊字符
hosts['Country'] = hosts['Host'].str.extract(r'(\b[A-Za-z]+(?:\s[A-Za-z]+)*\b)$')  # 提取主办国家
hosts['Country'] = hosts['Country'].map(noc_mapping)

"""
def preview_data():
    print("Medal Counts Dataset:\n", medal_counts.head())
    print("Hosts Dataset:\n", hosts.head())
    print("Programs Dataset:\n", programs.head())
    print("Athletes Dataset:\n", athletes.head())


preview_data()
"""

# ----------------------------------------------------- #


# 将国家名称转换为国家代码
medal_counts['NOC'] = medal_counts['NOC'].map(noc_mapping)

# 检查转换是否成功
# print("Medal Counts Dataset (After NOC Mapping):\n", medal_counts.head())

# 按年份和国家统计运动员数量
athlete_counts = athletes.groupby(['Year', 'NOC']).size().reset_index(name='AthleteCount')
# print("Athlete Counts Dataset:\n", athlete_counts.head())

# 将运动员数量合并到奖牌数据中
medal_counts = medal_counts.merge(athlete_counts, on=['Year', 'NOC'], how='left')
medal_counts['AthleteCount'] = medal_counts['AthleteCount'].fillna(0)  # 缺失值填充为0

# 检查合并结果
# print("Medal Counts Dataset (After Adding AthleteCount):\n", medal_counts.head())

# ----------------------------------------------------- #

# 将主办国家信息合并到奖牌数据中
medal_counts['Host'] = medal_counts['Year'].map(hosts.set_index('Year')['Country'].to_dict())
medal_counts.fillna(0, inplace=True)  # 填充缺失值为0

# ----------------------------------------------------- #

# 计算每个国家的历史奖牌平均值
historical_medals = medal_counts.groupby('NOC')[['Gold', 'Silver', 'Bronze', 'Total']].mean().reset_index()
historical_medals.rename(columns={
    'Gold': 'HistoricalGold',
    'Silver': 'HistoricalSilver',
    'Bronze': 'HistoricalBronze',
    'Total': 'HistoricalTotal'
}, inplace=True)

# 将历史奖牌数据合并到主数据中
medal_counts = medal_counts.merge(historical_medals, on='NOC', how='left')

# 添加是否为主办国的特征
medal_counts['IsHost'] = medal_counts['NOC'] == medal_counts['Host']

# ----------------------------------------------------- #

# 计算每年每个国家的男性和女性运动员数量
gender_counts = athletes.groupby(['Year', 'NOC', 'Sex']).size().unstack(fill_value=0)

# 计算男女运动员比例（男性人数/女性人数）
gender_counts['MaleToFemaleRatio'] = gender_counts['M'] / (gender_counts['F'] + 1e-5)  # 防止除以零

# 将男女运动员比例合并到奖牌数据中
medal_counts = medal_counts.merge(gender_counts[['MaleToFemaleRatio']], on=['Year', 'NOC'], how='left')

# 查看合并后的数据
print(medal_counts[['Year', 'NOC', 'MaleToFemaleRatio']].head())

# ----------------------------------------------------- #

# 确保数据按照年份排序
medal_counts = medal_counts.sort_values(by=['Year', 'NOC'])

# 使用shift来获取上届奖牌数据
medal_counts['PrevYearlyGold'] = medal_counts.groupby('NOC')['Gold'].shift(1)
medal_counts['PrevYearlySilver'] = medal_counts.groupby('NOC')['Silver'].shift(1)
medal_counts['PrevYearlyBronze'] = medal_counts.groupby('NOC')['Bronze'].shift(1)
medal_counts['PrevYearlyTotalMedals'] = medal_counts.groupby('NOC')[['Gold', 'Silver', 'Bronze']].shift(1).sum(axis=1)

# 填充缺失值（例如，第一届比赛的上届数据没有，填充为0）
medal_counts[['PrevYearlyGold', 'PrevYearlySilver', 'PrevYearlyBronze', 'PrevYearlyTotalMedals']] = medal_counts[
    ['PrevYearlyGold', 'PrevYearlySilver', 'PrevYearlyBronze', 'PrevYearlyTotalMedals']].fillna(0)

# 查看结果
print(medal_counts)

# ----------------------------------------------------- #


# ----------------------------------------------------- #

pd.set_option('display.max_columns', None)  # 显示所有列
print("Medal Counts Dataset:\n", medal_counts.head())

features = [
    'HistoricalGold', 'HistoricalSilver', 'HistoricalTotal',
    'IsHost', 'AthleteCount',
    'MaleToFemaleRatio',
]
labels_gold = medal_counts['Gold']  # 金牌数作为标签
X = medal_counts[features]

# 划分训练集和测试集
X_train, X_test, y_train_gold, y_test_gold = train_test_split(X, labels_gold, test_size=0.2, random_state=42)

# 模型训练
rf_gold = RandomForestRegressor(random_state=42)
rf_gold.fit(X_train, y_train_gold)

# 预测与误差计算
gold_predictions = rf_gold.predict(X_test)
residuals_gold = y_test_gold - gold_predictions  # 残差计算

# 图表绘制
fig, axes = plt.subplots(1, 2, figsize=(16, 6))

# 1. 实际值 vs 预测值对比图
sns.scatterplot(x=y_test_gold, y=gold_predictions, ax=axes[0], color='blue')
axes[0].plot([y_test_gold.min(), y_test_gold.max()], [y_test_gold.min(), y_test_gold.max()],
             color='red', linestyle='--', label='Ideal Prediction')  # 理想预测线
axes[0].set_xlabel('Actual Gold Medals')
axes[0].set_ylabel('Predicted Gold Medals')
axes[0].set_title('Actual vs Predicted Gold Medals')
axes[0].legend()

# 图像作用说明：此图展示了模型预测值与实际值的拟合程度。理想情况下，点应该接近红色虚线（y=x）。点越分散，说明预测误差越大。

# 2. 误差分布直方图
sns.histplot(residuals_gold, kde=True, ax=axes[1], color='green', bins=20)
axes[1].axvline(0, color='red', linestyle='--', label='Zero Error')  # 零误差线
axes[1].set_xlabel('Residuals (Actual - Predicted)')
axes[1].set_ylabel('Frequency')
axes[1].set_title('Residual Distribution')
axes[1].legend()

# 图像作用说明：此图反映了预测误差的分布情况。残差越集中在0附近，模型预测越准确。如果分布偏斜或范围过大，可能需要优化模型。

# 调整布局并显示图表
plt.tight_layout()
plt.show()

# 创建图表（去除残差图）
fig, axes = plt.subplots(1, 1, figsize=(8, 6))  # 只创建一个子图

# 特征重要性的条形图（Bar Plot）
importance_gold = rf_gold.feature_importances_
axes.bar(features, importance_gold, color='blue')
axes.set_xlabel('Features')
axes.set_ylabel('Importance')
axes.set_title('Feature Importance for Gold Medals Prediction')

# 设置水平坐标轴的字体大小
plt.xticks(rotation=45, fontsize=8)  # 旋转x轴标签，并调整字体大小为8

# 显示图表
plt.tight_layout()
plt.show()

# 输出误差指标
mse = mean_squared_error(y_test_gold, gold_predictions)
mae = mean_absolute_error(y_test_gold, gold_predictions)
print(f"Mean Squared Error (MSE): {mse:.2f}")
print(f"Mean Absolute Error (MAE): {mae:.2f}")
